import torch
from torch import Tensor

from .base import Metric


class Accuracy(Metric):
    """Metric that calculates accuracy.
    """

    def __init__(self, name: str = 'accuracy'):
        super(Accuracy, self).__init__(name)

    def update(self, pred: Tensor, label: Tensor):
        pred_label = torch.argmax(pred, dim=-1)
        label = label.long()

        num_correct = (pred_label == label).sum().item()
        self.num_inst += label.shape[0]
        self.sum_metric += num_correct
